/* Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * openUBMC is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

#include "context_config.h"
#include "spdm_command_impl.h"
#include "measurements/measurement_set.h"

MeasurementSet::MeasurementSet(const string& specification, const string& summaryType,
    const string& summaryHashAlgorithm)
    : specification(specification), summaryType(summaryType), summaryHashAlgorithm(summaryHashAlgorithm),
    summaryHash("") {}

int32_t MeasurementSet::GetMeasurementSet(void *spdm_context, uint8_t slot_id)
{
    // 不签名、仅返回当前测量记录的hash
    uint8_t request_attribute = 0;

    // 查询所有测量块整合的测量记录
    uint8_t measurement_operation = SPDM_GET_MEASUREMENTS_REQUEST_MEASUREMENT_OPERATION_ALL_MEASUREMENTS;

    uint8_t number_of_blocks = 0;
    uint8_t measurement_record[LIBSPDM_MAX_MEASUREMENT_RECORD_SIZE];
    uint32_t  measurement_record_length = sizeof(measurement_record);

    uint32_t status = ::GetMeasurement(spdm_context, nullptr, request_attribute, measurement_operation, slot_id,
        nullptr, &number_of_blocks, &measurement_record_length, measurement_record);
    if (LIBSPDM_STATUS_IS_ERROR(status)) {
        debug_log(DLOG_ERROR, "get measurement blocks num failed, status: %u", status);
        return RET_ERR;
    }

    this->summaryHash = ConvertBytesToString(measurement_record, measurement_record_length);

    return RET_OK;
}

int32_t MeasurementSet::GetSubMeasurements(void *spdm_context, const string& hash_algo, uint8_t slot_id)
{
    // 不签名、仅返回当前测量记录的hash
    uint8_t request_attribute = 0;

    // 查询可用的测量块总数
    uint8_t measurement_operation = SPDM_GET_MEASUREMENTS_REQUEST_MEASUREMENT_OPERATION_TOTAL_NUMBER_OF_MEASUREMENTS;

    uint8_t number_of_blocks = 0;

    uint32_t status = ::GetMeasurement(spdm_context, nullptr, request_attribute, measurement_operation, slot_id,
        nullptr, &number_of_blocks, nullptr, nullptr);
    if (LIBSPDM_STATUS_IS_ERROR(status)) {
        debug_log(DLOG_ERROR, "get measurement blocks num failed, status: %u", status);
        return RET_ERR;
    }
    uint8_t measurement_record[LIBSPDM_MAX_MEASUREMENT_RECORD_SIZE];
    uint32_t  measurement_record_length = sizeof(measurement_record);
    // 记录一个最大测量块数量，后续要用，避免被覆写掉
    this->measurementCount = number_of_blocks;
    uint8_t max_number_of_blocks = number_of_blocks;

    string sub_hash = "";

    // 循环获取每个测量块的测量记录
    for (uint8_t index = 1; index <= max_number_of_blocks; index++) {
        if (index == max_number_of_blocks) {
            // 在最后一次交互生产签名摘要(libspdm会进行验签)
            request_attribute = SPDM_GET_MEASUREMENTS_REQUEST_ATTRIBUTES_GENERATE_SIGNATURE;
        }

        status = ::GetMeasurement(spdm_context, nullptr, request_attribute, index, slot_id, nullptr, &number_of_blocks,
            &measurement_record_length, measurement_record);
        if (LIBSPDM_STATUS_IS_ERROR(status)) {
            debug_log(DLOG_ERROR, "get measurement failed, index: %u, status: %u", index, status);
            return RET_ERR;
        }

        sub_hash = ConvertBytesToString(measurement_record, measurement_record_length);
        measurements.emplace(index, make_unique<Measurement>(index, hash_algo, "FirmwareConfiguration", sub_hash,
            time(nullptr), true, ""));
    }

    return RET_OK;
}

static bool check_indices(uint8_t measurement_count, vector<uint8_t> indices)
{
    if (indices.size() > measurement_count) {
        debug_log(DLOG_ERROR, "invalid indices, indices count(%u), provided measurement count(%u)", indices.size(),
            measurement_count);
        return false;
    }

    uint8_t flags[MAX_MEASUREMENT_COUNT + 1] = {};
    for (uint8_t idx : indices) {
        if (flags[idx] == 0) {
            flags[idx] = 1;
        } else {
            debug_log(DLOG_ERROR, "invalid indices, index(%u) repeated", idx);
            return false;
        }
    }

    if (flags[MAX_MEASUREMENT_COUNT] != 0 && indices.size() > 1) {
        debug_log(DLOG_ERROR, "invalid indices, if contains 255 then can only contain asingle value of 255");
        return false;
    }

    return true;
}

static tuple<string, string, string> GetVersionAndAlgo(void *spdm_context)
{
    libspdm_data_parameter_t parameter = {};
    parameter.location = LIBSPDM_DATA_LOCATION_CONNECTION;

    uint16_t version;
    size_t   version_size = sizeof(version);
    SpdmGetData(spdm_context, LIBSPDM_DATA_SPDM_VERSION, &parameter, &version, &version_size);

    uint32_t hash_algo;
    size_t   hasl_algo_size = sizeof(hash_algo);
    SpdmGetData(spdm_context, LIBSPDM_DATA_BASE_HASH_ALGO, &parameter, &hash_algo, &hasl_algo_size);

    uint32_t asym_algo;
    size_t   asym_algo_size = sizeof(asym_algo);
    SpdmGetData(spdm_context, LIBSPDM_DATA_BASE_ASYM_ALGO, &parameter, &asym_algo, &asym_algo_size);

    return {GetCapabilityString("Version", version),
        GetCapabilityString("HashAlgo", hash_algo),
        GetCapabilityString("AsymAlgo", asym_algo)};
}

static tuple<int32_t, uint32_t, vector<uint8_t>> GetMeasurementsWithoutNonceSig(void *spdm_context, uint8_t index,
    uint8_t slot_id)
{
    // 不签名、仅返回当前测量记录的hash
    uint8_t request_attribute = 0;

    uint8_t number_of_blocks = 0;
    vector<uint8_t> measurement_record(LIBSPDM_MAX_MEASUREMENT_RECORD_SIZE);
    uint32_t  measurement_record_length = sizeof(measurement_record);

    uint32_t status = ::GetMeasurement(spdm_context, nullptr, request_attribute, index, slot_id, nullptr,
        &number_of_blocks, &measurement_record_length, measurement_record.data());
    if (LIBSPDM_STATUS_IS_ERROR(status)) {
        debug_log(DLOG_ERROR, "get measurement failed, index: %u, status: %u", index, status);
        return {RET_ERR, 0, {}};
    }

    return {RET_OK, measurement_record_length, measurement_record};
}

static tuple<int32_t, uint32_t, vector<uint8_t>> GetMeasurementsWithNonceSig(void *spdm_context, uint8_t index, uint8_t slot_id,
    const string& nonce)
{
    // 在最后一次交互生产签名摘要(libspdm会进行验签)
    uint8_t request_attribute = SPDM_GET_MEASUREMENTS_REQUEST_ATTRIBUTES_GENERATE_SIGNATURE;

    uint8_t number_of_blocks = 0;
    vector<uint8_t> measurement_record(LIBSPDM_MAX_MEASUREMENT_RECORD_SIZE);
    uint32_t  measurement_record_length = sizeof(measurement_record);

    vector<uint8_t> requester_nonce_in(SPDM_NONCE_SIZE);
    vector<uint8_t> requester_nonce(SPDM_NONCE_SIZE);
    vector<uint8_t> responder_nonce(SPDM_NONCE_SIZE);
    vector<uint8_t> opaque_data(SPDM_MAX_OPAQUE_DATA_SIZE);
    size_t opaque_data_size = opaque_data.size();

    vector<uint8_t> nonce_in(SPDM_NONCE_SIZE);
    ConvertStringToBytes(nonce.c_str(), nonce.length(), nonce_in.data(), nonce_in.size());

    for (uint8_t i = 0; i < SPDM_NONCE_SIZE; i++) {
        requester_nonce[i] = 0x00;
        responder_nonce[i] = 0x00;
        requester_nonce_in[i] = nonce_in[i];
    }

    uint32_t status = ::GetMeasurementEx(spdm_context, nullptr, request_attribute, index, slot_id, nullptr,
        &number_of_blocks, &measurement_record_length, measurement_record.data(), requester_nonce_in.data(),
        requester_nonce.data(), responder_nonce.data(), opaque_data.data(), &opaque_data_size);
    
    if (LIBSPDM_STATUS_IS_ERROR(status)) {
        debug_log(DLOG_ERROR, "get signed measurement failed, index: %u, status: %u", index, status);
        return {RET_ERR, 0, {}};
    }

    return {RET_OK, measurement_record_length, measurement_record};
}

tuple<string, string, string, string, string> MeasurementSet::GetSignedMeasurement(void *spdm_context,
    vector<uint8_t> indices, uint8_t slot_id, const string& nonce)
{
    vector<uint8_t> processIndices = indices;
    if (processIndices.empty()) {
        debug_log(DLOG_INFO, "No measurement indices, request all measurements");
        processIndices = {255};
    }

    if (!check_indices(measurementCount, processIndices)) {
        debug_log(DLOG_ERROR, "check indices failed");
        throw runtime_error("check indices failed");
    }
    
    vector<uint8_t> all_measurements = {};
    uint8_t retrieved_measurement_count = 0;

    int32_t res;
    uint32_t len;
    vector<uint8_t> record;

    tuple<int32_t, uint32_t, vector<uint8_t>> result;

    for (uint8_t i = 0; i < indices.size(); i++) {
        if (i != indices.size() - 1) {
            result = GetMeasurementsWithoutNonceSig(spdm_context, indices[i], slot_id);
        } else {
            result = GetMeasurementsWithNonceSig(spdm_context, indices[i], slot_id, nonce);
        }

        tie(res, len, record) = result;

        if (res != RET_OK) {
            break;
        }

        record.resize(len);

        all_measurements.insert(all_measurements.end(), record.begin(), record.end());
        retrieved_measurement_count += 1;

        debug_log(DLOG_INFO, "retrieved measurement %u, size: %u", indices[i], len);
    }

    if (retrieved_measurement_count != indices.size()) {
        debug_log(DLOG_ERROR, "can not retrieved all blocks. retrieved: %u, required: %u",
            retrieved_measurement_count, indices.size());
        throw runtime_error("get signed measurements failed");
    }

    string base64_measurement = base64Encode(all_measurements);

    auto [version, hash_algo, asym_algo] = GetVersionAndAlgo(spdm_context);

    return {version, "", hash_algo, asym_algo, base64_measurement};
}